import hmac
import logging
import struct
from collections import defaultdict
from random import random, seed
from time import time_ns

from scapy.compat import raw
from scapy.layers.inet import IP, TCP, UDP

from policy import *
import sys

sys.path.append('/home/nfv/vSFC-p')
from NFs.network_address_translation import NAT
from NFs.merger import Merger
from NFs.intrusion_detection_system import IDS
from NFs.firewall import Firewall


# 包正确性校验协议
class PacketCorrectness(object):
    def __init__(self, network_function, network_function_identifier: int, name: str):
        self.network_function = network_function  #
        self.network_function_identifier = network_function_identifier  # definitions in policy.py
        self.name = name

        self.in_dst_nf = None
        self.in_src_nf = None
        self.in_pkt_identifier = None
        self.in_auth = None

        self.out_dst_nf = None
        self.out_src_nf = None
        self.out_pkt_identifier = None
        self.out_auth = None

        self.key = b'packet'
        self.pkt_id_count = 0
        # 采样
        self.is_sampled_packet = False
        self.per_flow_sample_para = defaultdict(float)  # 每个流的采样率

        # TODO：logging, may be async
        _format = '%(asctime)s - %(levelname)s - %(filename)s [line:%(lineno)d] : %(message)s'
        logging.basicConfig(format=_format, filename=self.name, level=logging.INFO)
        self.logger = logging.getLogger(self.name)

        self.original_ttl = 64
        seed(1)

    def un_pack(self, pkt: bytes):
        # 解析出必须的字段
        self.in_pkt_identifier, self.in_src_nf, self.in_dst_nf = struct.unpack('>III', pkt[-28:-16])
        self.in_auth = pkt[-16:]

    def pack(self):
        res = struct.pack('>III', self.out_pkt_identifier, self.out_src_nf, self.out_dst_nf)
        return res

    def decrease_len_chksum(self, pkt: bytes):
        # 重新生成chksum, 更新len
        scapy_pkt = IP(pkt)
        del scapy_pkt[IP].chksum
        scapy_pkt[IP].len -= 28  # 还原最初的包
        # 不校验ttl
        self.original_ttl = scapy_pkt[IP].ttl
        scapy_pkt[IP].ttl = 64
        return bytes(IP(bytes(scapy_pkt)))

    def increase_len_chksum(self, pkt: bytes):
        scapy_pkt = IP(pkt)
        del scapy_pkt[IP].chksum
        scapy_pkt[IP].len += 28
        scapy_pkt[IP].ttl = self.original_ttl
        return bytes(IP(bytes(scapy_pkt)))

    def retain_len_chksum(self, pkt: bytes):
        scapy_pkt = IP(pkt[:-28])
        del scapy_pkt[IP].chksum
        scapy_pkt[IP].ttl = self.original_ttl
        return bytes(IP(bytes(scapy_pkt)))

    def verify_integrity(self, pkt: bytes):  # 校验
        self.un_pack(pkt)  # 解析出字段
        self.is_sampled_packet = False
        # if self.network_function_identifier != self.in_dst_nf:
        if self.network_function_identifier not in self.policy(self.in_src_nf):  # 后一跳有多个时
            self.logger.error(f"Invalid destination {self.in_dst_nf} for {self.network_function_identifier}.")
            return False
        msg = pkt[:-16]  # pkt|pktID|srcNF|dstNF
        auth = hmac.new(self.key, msg + b'0', digestmod='MD5').digest()
        if auth != self.in_auth:
            auth = hmac.new(self.key, msg + b'1', digestmod='MD5').digest()
            if auth != self.in_auth:
                self.logger.error(f"Invalid packet {self.in_pkt_identifier, self.in_dst_nf}, may be injected.")
                return False
            else:
                self.is_sampled_packet = True
                return True
        else:
            return True

    def generate(self, pkt: bytes):  # 网关处
        flow_hash = self.compute_flow_hash(pkt)
        self.out_src_nf = self.network_function_identifier
        self.out_pkt_identifier = self.generate_pkt_identifier()
        self.out_dst_nf = self.policy(self.out_src_nf)[0]
        msg = pkt + self.pack()
        if self.sample(flow_hash):
            self.out_auth = hmac.new(self.key, msg + b'1', digestmod='MD5').digest()
            self.logger.info((b'null', msg + self.out_auth, self.network_function_identifier, time_ns()))
        else:
            self.out_auth = hmac.new(self.key, msg + b'0', digestmod='MD5').digest()
        return self.increase_len_chksum(msg + self.out_auth)

    @staticmethod
    def compute_flow_hash(pkt: bytes) -> int:
        scapy_pkt = IP(pkt)
        src, dst = scapy_pkt[IP].src, scapy_pkt[IP].dst
        proto = scapy_pkt[IP].proto
        if proto == 6:  # TCP
            src_port, dst_port = scapy_pkt[TCP].sport, scapy_pkt[TCP].dport
        else:
            src_port, dst_port = scapy_pkt[UDP].sport, scapy_pkt[UDP].dport
        flow_hash = hash(src + dst + str(proto) + str(src_port) + str(dst_port))
        return flow_hash

    def sample(self, flow_hash, alpha=0.1, min_val=0.1):
        if flow_hash not in self.per_flow_sample_para:
            self.per_flow_sample_para[flow_hash] = 1.0
        sample_ratio = min_val if self.per_flow_sample_para[flow_hash] < min_val else self.per_flow_sample_para[
            flow_hash]
        r = random()
        if r < sample_ratio:
            if sample_ratio > self.per_flow_sample_para[flow_hash]:  # 取最小值了
                return True
            else:
                self.per_flow_sample_para[flow_hash] /= (1 + alpha)  # 抵制计算
                return True
        return False

    def process(self, pkt: bytes):
        pkt = self.decrease_len_chksum(pkt)  # 还原最初的包
        if self.verify_integrity(pkt):
            msg = pkt[:-28]  # pkt
            out_msg = self.network_function.match(msg)
            if out_msg is False:
                return
            elif out_msg is True:  # 对于IDS和FW
                out_msg = msg
            self.out_pkt_identifier = self.in_pkt_identifier
            self.out_src_nf = self.network_function_identifier
            self.out_dst_nf = self.policy(self.out_src_nf)[0]  # 只有一个
            msg = out_msg + self.pack()
            if self.is_sampled_packet:
                self.out_auth = hmac.new(self.key, msg + b'1', digestmod='MD5').digest()
                self.logger.info((pkt, msg + self.out_auth, self.network_function_identifier, time_ns()))
            else:
                self.out_auth = hmac.new(self.key, msg + b'0', digestmod='MD5').digest()
            return self.increase_len_chksum(msg + self.out_auth)

    def merge(self, pkt: bytes):
        # TODO：队列中需要特别对待
        pkt = self.decrease_len_chksum(pkt)  # 还原最初的包
        if self.verify_integrity(pkt):
            out_msg, data = self.network_function.match(pkt, self.in_pkt_identifier, self.in_src_nf)  # 合并结果
            if out_msg is False:
                return False  # 输入条件未满足
            else:
                out_msg = out_msg[:-28]
            # 返回了合并结果
            self.out_pkt_identifier = self.in_pkt_identifier
            self.out_src_nf = self.network_function_identifier  # 具有相同的pktID
            dst_nfs = self.policy(self.out_src_nf)
            inputs, outputs = [], []
            for packet, _ in data:
                inputs.append(self.increase_len_chksum(packet))
            for dst_nf in dst_nfs:
                self.out_dst_nf = dst_nf
                msg = out_msg + self.pack()
                if self.is_sampled_packet:
                    self.out_auth = hmac.new(self.key, msg + b'1', digestmod='MD5').digest()
                    out_pkt = msg + self.out_auth
                    outputs.append(self.increase_len_chksum(out_pkt))
                    for packet in inputs:
                        self.logger.info((packet, out_pkt, self.network_function_identifier, time_ns()))
                else:
                    self.out_auth = hmac.new(self.key, msg + b'0', digestmod='MD5').digest()
                    outputs.append(self.increase_len_chksum(msg + self.out_auth))
            return outputs

    def exit_gateway(self, pkt: bytes):
        pkt = self.decrease_len_chksum(pkt)  # 还原最初的包
        if self.verify_integrity(pkt):
            if self.is_sampled_packet:
                self.logger.info((pkt, b'null', self.network_function_identifier, time_ns()))
            return self.retain_len_chksum(pkt)

    def generate_pkt_identifier(self):
        self.pkt_id_count += 1
        if self.pkt_id_count == 4294967295:  # 2^32-1
            self.pkt_id_count = 0
        return self.pkt_id_count

    @staticmethod
    def policy(src_nf: int, pkt_identifier: int = 0, flow_identifier: int = 0):
        # next-hop NF
        # not use pkt_identifier, as all packet go through same path in the experiment
        return policies[src_nf]['out']


if __name__ == '__main__':
    packet_correctness = PacketCorrectness(None, 1, 'GatewayIn')
    _pkt = b'E\x00\x00<\x00\x01\x00\x00@\x06\x11\xe5\xf7\x89\xe9\xcc$\x91b\xef\x8duM\xf1\x00\x00\x00\x00\x00\x00\x00' \
           b'\x00P\x02 \x00!\xca\x00\x00abcdefffffffffffffff'
    start = time_ns()
    # 测试网关
    o_pkt = None
    for i in range(100000):
        o_pkt = packet_correctness.generate(_pkt)
        break
        # packet = packet_correctness.exit_gateway(out_pkt)
    print((time_ns() - start) / 10 ** 9, "GatewayIn")
    print(o_pkt)

    # 测试NAT
    nat = NAT("172.10.1.2", 8080)
    packet_correctness_nat = PacketCorrectness(nat, 10, 'NAT')
    start = time_ns()
    for i in range(100000):
        o_pkt = packet_correctness_nat.process(o_pkt)
        break
        # print(packet_correctness_nat.sample_flag)
    print((time_ns() - start) / 10 ** 9, "NAT")
    print(o_pkt)

    # 测试merger
    merger = Merger(1, 2)
    packet_correctness_merger = PacketCorrectness(merger, 11, 'Merger-1-2')
    outs = None
    start = time_ns()
    for i in range(100000):
        outs = packet_correctness_merger.merge(o_pkt)
        break
    print((time_ns() - start) / 10 ** 9, "Merger-1-2")

    print(outs[0])
    print(outs[1])
    # 测试IDS
    ids = IDS('../NFs/ids_rules.csv')
    packet_correctness_ids = PacketCorrectness(ids, 110, 'IDS')
    start = time_ns()
    ids_o, fw_o = None, None
    for i in range(100000):
        ids_o = packet_correctness_ids.process(outs[0])
        break
    print((time_ns() - start) / 10 ** 9, "IDS")
    print(ids_o)

    # 测试FW
    fw = Firewall()
    packet_correctness_fw = PacketCorrectness(fw, 1110, 'FW')
    start = time_ns()
    for i in range(100000):
        fw_o = packet_correctness_fw.process(outs[1])
        break
    print((time_ns() - start) / 10 ** 9, "FW")

    # 测试merger
    merger = Merger(1, 1)  # 2,1
    packet_correctness_merger = PacketCorrectness(merger, 11110, 'Merger-2-1')
    outs = None
    start = time_ns()
    for i in range(100000):
        outs = packet_correctness_merger.merge(ids_o)
        # outs = packet_correctness_merger.merge(fw_o)
        break
    print((time_ns() - start) / 10 ** 9, "Merger-2-1")
    print(outs[0])

    # 测试出网关
    packet_correctness = PacketCorrectness(None, 0, 'GatewayOut')
    start = time_ns()
    # 测试网关
    for i in range(100000):
        o_pkt = packet_correctness.exit_gateway(outs[0])
        break
        # packet = packet_correctness.exit_gateway(out_pkt)
    print((time_ns() - start) / 10 ** 9, "GatewayOut")
    print(o_pkt)
